open Sast

module NameMap = Map.Make(struct
  type t = string
  let compare x y = Pervasives.compare x y
end)

exception ReturnException of int
exception Bug of string (* For "impossible" situations *)

(* Main entry point: run a program *)
let run prog args =
    (* Find and return symbols from scope *)
    let rec getvar id = function
	[] -> raise (Failure ("undeclared identifier " ^ id))
	| loc::scp ->
	    let (vars,funs) = loc in
	    if NameMap.mem id vars then
		NameMap.find id vars
	    else
		getvar id scp
    in
    let rec getfun id = function
	[] -> raise (Failure ("undefined function " ^ id))
	| loc::scp ->
	    let (vars,funs) = loc in
	    if NameMap.mem id funs then
		NameMap.find id funs
	    else
		getfun id scp
    in

    (* Evaluate an expression and return value *)
    let rec eval scope = function
      Literal(i) -> i
    | Noexpr -> 1 (* must be non-zero for the for loop predicate *)
    | Id(var) -> getvar var scope
    | Uop (op, e) ->
        let v = eval scope e in
        let boolean i = if i then 1 else 0 in
        (match uop with
          Not -> boolean(!v)
        | Minus -> -v
    | Binop(e1, op, e2) ->
	let v1 = eval scope e1 in
	let v2 = eval scope e2 in
	let boolean i = if i then 1 else 0 in
	(match op with
	  Add -> v1 + v2
	| Sub -> v1 - v2
	| Mult -> v1 * v2
	| Div -> v1 / v2
	| Equal -> boolean (v1 = v2)
	| Neq -> boolean (v1 != v2)
	| Less -> boolean (v1 < v2)
	| Leq -> boolean (v1 <= v2)
	| Greater -> boolean (v1 > v2)
	| Geq -> boolean (v1 >= v2))
    | Call(f, actuals) ->
	let fdecl = getfun f scope in
	let actuals = List.fold_left
	    (fun actuals actual ->
		let v = eval scope actual in v :: actuals) [] actuals
	in
	try call fdecl actuals scope
	with ReturnException(v) -> v
    
    (* Invoke a function and return value *)
    and call fdecl actuals globals =
	(* Enter the function: bind actual values to formal arguments *)
	let (fform, fbody) = fdecl in
	let lvars =
	    try List.fold_left2
		(fun locals formal actual -> 
		    let _, name = formal in 
		    NameMap.add name actual locals)
		NameMap.empty fform actuals
	    with Invalid_argument(_) ->
		raise (Failure ("wrong number of arguments passed to function"))
	in
	(* Execute function body, return returned value, ignore scope *)
	let ret, _ = exec ((lvars,NameMap.empty)::globals) fbody in ret

    (* Run through a list of statements.  Return value of last statement and
     * modified scope.  This is used in enough places to justify a separate
     * function for it. *)
    and stmtBlock scp statlist = 
        List.fold_left (fun ret stmt -> 
            let _ ,scp' = ret in exec scp' stmt) (0, scp) statlist

    (* Execute a statement and return a value and updated scope *)
    and exec scope = function
    Block(stmts) -> 
        let ret, _ = stmtBlock scope stmts in ret, scope
    | Expr(e) -> eval scope e, scope
    | FDclr(tp, lname, formals, body) -> (match scope with
        [] -> raise (Bug ("empty scope in FDclr"))
        | (lvars, lfuns)::globals -> 
            let lfuns = NameMap.add lname (formals,body) lfuns in
            0, ((lvars, lfuns)::globals))
    | VDclr(tp, lname, actual) -> (match scope with
        [] -> raise (Bug ("empty scope in VDclr"))
        | (lvars, lfuns)::globals -> 
            let value, _ = exec scope actual in
            let lvars = NameMap.add lname value lvars in
            value, (lvars, lfuns)::globals)
    | If(cond, truebody, falsebody) -> 
            let ret, _ =
                if (eval scope cond) = 0
                then exec scope falsebody
                else exec scope truebody
            in ret, scope
    | Scope(init, body) ->  (* first run a statement block, then body *)
            let _, scope' = stmtBlock scope init in
            let ret, _ = exec scope' body in
            ret, scope
    | While(body, cond) ->
            let rec loop scp = 
                let value, scp' = stmtBlock scp body in
                match eval scp' cond with
                0 -> value, scope
                | _ -> loop scp'
            in loop scope
    in

    (* Run a program: start with an empty scope and run through program.
     * Then find and call "begin", and print it's result *)
    let _, scope = stmtBlock [(NameMap.empty,NameMap.empty)] prog in
    try call (getfun "begin" scope) args scope
    with Failure(s) -> raise (Failure s)(*("did not find the begin() function"))*)

